#!/usr/bin/env python3
"""
种子重放工具
从文件中读取种子，并通过socket发送到指定的目标服务器
"""

import argparse
import socket
import sys
import time
import Fuzzer

def read_test_case_from_file(file_path):
    """
    从文件中读取test case，解析保存的消息格式
    
    Args:
        file_path: 输入文件路径
        
    Returns:
        List[bytes]: 消息列表，每个元素是bytes类型的协议消息
    """
    messages = []
    
    try:
        with open(file_path, 'rb') as f:
            while True:
                # 读取4字节长度（小端序）
                length_bytes = f.read(4)
                if len(length_bytes) < 4:
                    break  # 文件结束
                
                length = int.from_bytes(length_bytes, byteorder='little', signed=False)
                
                # 读取对应长度的消息数据
                message = f.read(length)
                if len(message) < length:
                    break  # 文件意外结束
                
                messages.append(bytes(message))
                
    except (IOError, OSError) as e:
        print(f"Error reading test case from {file_path}: {e}", file=sys.stderr)
        return []
    
    return messages


def send_messages(host, port, messages, timeout=5.0, receive_response=False):
    """
    通过socket连接发送消息到目标服务器
    
    Args:
        host: 目标服务器IP地址
        port: 目标服务器端口
        messages: 要发送的消息列表，每个元素是bytes类型
        timeout: socket超时时间（秒）
        receive_response: 是否接收并打印响应
        
    Returns:
        bool: 是否成功发送
    """
    if not messages:
        print("No messages to send", file=sys.stderr)
        return False
    
    sock = None
    try:
        # 创建socket连接
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        
        # 连接到目标服务器
        print(f"Connecting to {host}:{port}...")
        sock.connect((host, port))
        print(f"Connected successfully")
        
        # 发送所有消息
        print(f"Sending {len(messages)} message(s)...")
        for i, msg in enumerate(messages, 1):
            if isinstance(msg, (list, tuple)):
                # 如果消息是列表或元组，转换为bytes
                msg_bytes = bytes(msg)
            else:
                # 确保是bytes类型
                msg_bytes = bytes(msg) if not isinstance(msg, bytes) else msg
            
            # 发送消息
            sent = sock.sendall(msg_bytes)
            print(f"  [Message {i}/{len(messages)}] Sent {len(msg_bytes)} bytes")
            time.sleep(0.1)  # 短暂延迟，避免消息过快
        
        print("All messages sent successfully")
        
        # 如果要求接收响应
        if receive_response:
            print("\nWaiting for response...")
            try:
                response = sock.recv(4096)
                if response:
                    print(f"Received {len(response)} bytes:")
                    print(f"  Hex: {response.hex()}")
                    print(f"  Preview: {response[:100]}...")
                else:
                    print("No response received")
            except socket.timeout:
                print("Timeout waiting for response")
            except Exception as e:
                print(f"Error receiving response: {e}")
        
        return True
        
    except socket.timeout:
        print(f"Timeout connecting to {host}:{port}", file=sys.stderr)
        return False
    except socket.error as e:
        print(f"Socket error: {e}", file=sys.stderr)
        return False
    except Exception as e:
        print(f"Unexpected error: {e}", file=sys.stderr)
        return False
    finally:
        if sock:
            sock.close()
            print("Connection closed")


def main():
    parser = argparse.ArgumentParser(
        description='重放种子文件：读取种子并通过socket发送到目标服务器',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  %(prog)s seed.bin 127.0.0.1 4433
  %(prog)s seed.bin localhost 4433 --timeout 10 --receive-response
        """
    )
    
    parser.add_argument('seed_file', 
                       help='种子文件路径')
    parser.add_argument('host', 
                       help='目标服务器IP地址或主机名')
    parser.add_argument('port', 
                       type=int,
                       help='目标服务器端口')
    parser.add_argument('--timeout', '-t',
                       type=float,
                       default=5.0,
                       help='Socket超时时间（秒），默认: 5.0')
    parser.add_argument('--receive-response', '-r',
                       action='store_true',
                       help='接收并显示服务器响应')
    
    args = parser.parse_args()
    
    # 读取种子文件
    print(f"Reading seed file: {args.seed_file}")
    messages = read_test_case_from_file(args.seed_file)
    
    if not messages:
        print("Failed to read messages from seed file", file=sys.stderr)
        sys.exit(1)
    
    print(f"Loaded {len(messages)} message(s) from seed file\n")
    
    # 发送消息
    success = send_messages(args.host, args.port, messages, 
                           timeout=args.timeout,
                           receive_response=args.receive_response)
    
    if success:
        print("\nReplay completed successfully")
        sys.exit(0)
    else:
        print("\nReplay failed", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()

